{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Maximum Elegance of a K-Length Subsequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #greedy #array #hash-table #sorting #heap-priority-queue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #贪心 #数组 #哈希表 #排序 #堆（优先队列）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: findMaximumElegance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #子序列最大优雅度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个长度为 <code>n</code> 的二维整数数组 <code>items</code> 和一个整数 <code>k</code> 。</p>\n",
    "\n",
    "<p><code>items[i] = [profit<sub>i</sub>, category<sub>i</sub>]</code>，其中 <code>profit<sub>i</sub></code> 和 <code>category<sub>i</sub></code> 分别表示第 <code>i</code> 个项目的利润和类别。</p>\n",
    "\n",
    "<p>现定义&nbsp;<code>items</code> 的 <strong>子序列</strong> 的 <strong>优雅度</strong> 可以用 <code>total_profit + distinct_categories<sup>2</sup></code> 计算，其中 <code>total_profit</code> 是子序列中所有项目的利润总和，<code>distinct_categories</code> 是所选子序列所含的所有类别中不同类别的数量。</p>\n",
    "\n",
    "<p>你的任务是从 <code>items</code> 所有长度为 <code>k</code> 的子序列中，找出 <strong>最大优雅度</strong> 。</p>\n",
    "\n",
    "<p>用整数形式表示并返回 <code>items</code> 中所有长度恰好为 <code>k</code> 的子序列的最大优雅度。</p>\n",
    "\n",
    "<p><strong>注意：</strong>数组的子序列是经由原数组删除一些元素（可能不删除）而产生的新数组，且删除不改变其余元素相对顺序。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>items = [[3,2],[5,1],[10,1]], k = 2\n",
    "<strong>输出：</strong>17\n",
    "<strong>解释：\n",
    "</strong>在这个例子中，我们需要选出长度为 2 的子序列。\n",
    "其中一种方案是 items[0] = [3,2] 和 items[2] = [10,1] 。\n",
    "子序列的总利润为 3 + 10 = 13 ，子序列包含 2 种不同类别 [2,1] 。\n",
    "因此，优雅度为 13 + 2<sup>2</sup> = 17 ，可以证明 17 是可以获得的最大优雅度。 \n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>items = [[3,1],[3,1],[2,2],[5,3]], k = 3\n",
    "<strong>输出：</strong>19\n",
    "<strong>解释：</strong>\n",
    "在这个例子中，我们需要选出长度为 3 的子序列。 \n",
    "其中一种方案是 items[0] = [3,1] ，items[2] = [2,2] 和 items[3] = [5,3] 。\n",
    "子序列的总利润为 3 + 2 + 5 = 10 ，子序列包含 3 种不同类别 [1, 2, 3] 。 \n",
    "因此，优雅度为 10 + 3<sup>2</sup> = 19 ，可以证明 19 是可以获得的最大优雅度。</pre>\n",
    "\n",
    "<p><strong>示例 3：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>items = [[1,1],[2,1],[3,1]], k = 3\n",
    "<strong>输出：</strong>7\n",
    "<strong>解释：\n",
    "</strong>在这个例子中，我们需要选出长度为 3 的子序列。\n",
    "我们需要选中所有项目。\n",
    "子序列的总利润为 1 + 2 + 3 = 6，子序列包含 1 种不同类别 [1] 。\n",
    "因此，最大优雅度为 6 + 1<sup>2</sup> = 7 。</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= items.length == n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>items[i].length == 2</code></li>\n",
    "\t<li><code>items[i][0] == profit<sub>i</sub></code></li>\n",
    "\t<li><code>items[i][1] == category<sub>i</sub></code></li>\n",
    "\t<li><code>1 &lt;= profit<sub>i</sub> &lt;= 10<sup>9</sup></code></li>\n",
    "\t<li><code>1 &lt;= category<sub>i</sub> &lt;= n </code></li>\n",
    "\t<li><code>1 &lt;= k &lt;= n</code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [maximum-elegance-of-a-k-length-subsequence](https://leetcode.cn/problems/maximum-elegance-of-a-k-length-subsequence/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [maximum-elegance-of-a-k-length-subsequence](https://leetcode.cn/problems/maximum-elegance-of-a-k-length-subsequence/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[[3,2],[5,1],[10,1]]\\n2', '[[3,1],[3,1],[2,2],[5,3]]\\n3', '[[1,1],[2,1],[3,1]]\\n3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x : -x[0]) \n",
    "        ans = tot = 0 \n",
    "        vis = set()\n",
    "        dupl = []\n",
    "        for i, (p, c) in enumerate(items[:k]):\n",
    "            tot += p \n",
    "            if c in vis:\n",
    "                dupl.append(p) \n",
    "            vis.add(c) \n",
    "        ans = tot + len(vis) * len(vis) \n",
    "\n",
    "        for i, (p, c) in enumerate(items[k:]):\n",
    "            if c not in vis:\n",
    "                if not dupl:\n",
    "                   break \n",
    "                #print(dupl)\n",
    "                pre_p = dupl.pop() \n",
    "                tot += p - pre_p \n",
    "                vis.add(c)\n",
    "                ans = max(ans, tot + len(vis) * len(vis)) \n",
    "        return ans  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, s: List[List[int]], k: int) -> int:\n",
    "        n = len(s)\n",
    "        s.sort(reverse=True)\n",
    "        st = set()\n",
    "        f = [False] * k\n",
    "        for i in range(k):\n",
    "            if s[i][1] in st:\n",
    "                continue\n",
    "            st.add(s[i][1])\n",
    "            f[i] = True\n",
    "        m = len(st)\n",
    "        ans = t = sum((v for v, _ in s[: k])) + m * m\n",
    "        i = k - 1\n",
    "        j = k\n",
    "        while 1:\n",
    "            while i >= 0 and f[i]:\n",
    "                i -= 1\n",
    "            if i < 0:\n",
    "                break\n",
    "            while j < n and s[j][1] in st:\n",
    "                j += 1\n",
    "            if j >= n:\n",
    "                break\n",
    "            st.add(s[j][1])\n",
    "            t += s[j][0] - s[i][0]\n",
    "            t += 2 * m + 1\n",
    "            m += 1\n",
    "            ans = max(ans, t)\n",
    "            i -= 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key = lambda i:-i[0])\n",
    "        ans = total_profit = 0\n",
    "        vis = set()\n",
    "        duplicate = []\n",
    "        for i ,(profit,category) in enumerate(items):\n",
    "            if i < k:\n",
    "                total_profit += profit\n",
    "                if category not in vis:\n",
    "                    vis.add(category)\n",
    "                else:\n",
    "                    duplicate.append(profit)\n",
    "            elif duplicate and category not in vis:\n",
    "                vis.add(category)\n",
    "                total_profit += profit - duplicate.pop()\n",
    "            ans = max(ans,total_profit + len(vis) * len(vis))\n",
    "        return ans \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda i: -i[0])  # 把利润从大到小排序\n",
    "        ans = total_profit = 0\n",
    "        vis = set()\n",
    "        duplicate = []  # 重复类别的利润\n",
    "        for i, (profit, category) in enumerate(items):\n",
    "            if i < k:\n",
    "                total_profit += profit\n",
    "                if category not in vis:\n",
    "                    vis.add(category)\n",
    "                else:  # 重复类别\n",
    "                    duplicate.append(profit)\n",
    "            elif duplicate and category not in vis:\n",
    "                vis.add(category)\n",
    "                total_profit += profit - duplicate.pop()  # 用最小利润替换\n",
    "            # else: 比前面的利润小，而且类别还重复了，选它只会让 total_profit 变小，len(vis) 不变，优雅度不会变大\n",
    "            ans = max(ans, total_profit + len(vis) * len(vis))\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda i: -i[0])  # 把利润从大到小排序\n",
    "        ans = total_profit = 0\n",
    "        vis = set()\n",
    "        duplicate = []  # 重复类别的利润\n",
    "        for i, (profit, category) in enumerate(items):\n",
    "            if i < k:\n",
    "                total_profit += profit\n",
    "                if category not in vis:\n",
    "                    vis.add(category)\n",
    "                else:  # 重复类别\n",
    "                    duplicate.append(profit)\n",
    "            elif duplicate and category not in vis:\n",
    "                vis.add(category)\n",
    "                total_profit += profit - duplicate.pop()  # 用最小利润替换\n",
    "            # else: 比前面的利润小，而且类别还重复了，选它只会让 total_profit 变小，len(vis) 不变，优雅度不会变大\n",
    "            ans = max(ans, total_profit + len(vis) * len(vis))\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key = lambda i:-i[0])\n",
    "        n = len(items)\n",
    "        total = 0\n",
    "        visited = set()\n",
    "        duplicate = []\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            if i < k:\n",
    "                total += items[i][0]\n",
    "                if items[i][1] not in visited:\n",
    "                    visited.add(items[i][1])\n",
    "                else:\n",
    "                    duplicate.append(items[i][0]) #比第一个出现的小\n",
    "            elif duplicate and items[i][1] not in visited:\n",
    "                total += items[i][0]-duplicate.pop()\n",
    "                visited.add(items[i][1])\n",
    "            ans = max(ans, total+len(visited)*len(visited))\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "      n = len(items)\n",
    "      items.sort(key=lambda x:-x[0])\n",
    "\n",
    "      # 前k个最优，并且记录哪些类别是重复出现\n",
    "      vis = set()\n",
    "      duplicate = list()\n",
    "      category_cnt = 0\n",
    "      res = 0\n",
    "      for profit, category in items[:k]:\n",
    "        if category not in vis:\n",
    "          vis.add(category)\n",
    "          category_cnt+=1\n",
    "        else:\n",
    "          duplicate.append((profit, category))\n",
    "        res += profit\n",
    "      tmp_res = res\n",
    "      res += category_cnt ** 2\n",
    "      # 从k+1开始，对于新类别，加入，并弹出一个重复的旧类别；对于已有类别，跳过\n",
    "      for profit, category in items[k:]:\n",
    "        if category not in vis and len(duplicate)>0:\n",
    "          vis.add(category)\n",
    "          old_pro, old_cate = duplicate.pop()\n",
    "          category_cnt += 1\n",
    "          tmp_res = tmp_res - old_pro + profit\n",
    "          res = max(res, tmp_res + category_cnt**2)\n",
    "      return res\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "\n",
    "        items.sort(reverse=True)\n",
    "\n",
    "        s = set()\n",
    "        dups = []\n",
    "        el = 0\n",
    "\n",
    "        for i in range(k):\n",
    "            el += items[i][0]\n",
    "            if items[i][1] not in s:\n",
    "                s.add(items[i][1])\n",
    "            else:\n",
    "                dups.append(i)\n",
    "\n",
    "        el += len(s) * len(s)\n",
    "\n",
    "        ret = el\n",
    "\n",
    "        for i in range(k, n):\n",
    "            if items[i][1] not in s:\n",
    "                if dups:\n",
    "                    di = dups.pop()\n",
    "                    el -= items[di][0] + len(s) * len(s)\n",
    "                    s.add(items[i][1])\n",
    "                    el += items[i][0] + len(s) * len(s)\n",
    "                    ret = max(ret, el)\n",
    "\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x:-x[0])\n",
    "        ans = s = 0\n",
    "        vis = set()\n",
    "        dup = []\n",
    "        for i, (p, t) in enumerate(items):\n",
    "            if i < k:\n",
    "                s += p\n",
    "                if t not in vis:\n",
    "                    vis.add(t)\n",
    "                else:\n",
    "                    dup.append(p)\n",
    "            else:\n",
    "                if dup and t not in vis:\n",
    "                    vis.add(t)\n",
    "                    s = s + p - dup.pop()\n",
    "            ans = max(ans, s + len(vis) * len(vis))\n",
    "                    \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        '''\n",
    "        排序\n",
    "        贪心，先保存profit大的\n",
    "        记录最小的（用栈）\n",
    "        记录种类（set）\n",
    "        '''\n",
    "        n=len(items)\n",
    "        res=tot=0\n",
    "        stack=[]\n",
    "        mset=set()\n",
    "        items.sort(key=lambda x: -x[0])\n",
    "        print(items)\n",
    "        for i in range(n):\n",
    "            if i<k:\n",
    "                #print(res,len(mset),len(stack))\n",
    "            \n",
    "                tot +=items[i][0]\n",
    "                if items[i][1] in mset:\n",
    "                    stack.append(items[i][0])\n",
    "                else: mset.add(items[i][1])\n",
    "            elif items[i][1] not in mset and stack:\n",
    "                    tot +=items[i][0]\n",
    "                    tot -=stack[-1]\n",
    "                    stack.pop()\n",
    "                    mset.add(items[i][1])\n",
    "            res=max(res,tot+len(mset)**2)    \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        cnt = [0]*(100010)\n",
    "        items.sort(key = lambda x:-x[0])\n",
    "        res = 0\n",
    "        h = []\n",
    "        nowtype = 0\n",
    "        print(items)\n",
    "        for i in range(k):\n",
    "            cost,types = items[i]\n",
    "            res += cost\n",
    "            if cnt[types]==0:nowtype+=1\n",
    "            cnt[types]+=1\n",
    "            if cnt[types]>1: h.append(cost)\n",
    "\n",
    "        res+=nowtype**2\n",
    "        heapify(h)\n",
    "        print(h,res,nowtype)\n",
    "        ans = res\n",
    "\n",
    "        for i in range(k,n):\n",
    "            cost,types = items[i]\n",
    "            if cnt[types]>0:continue   #如果这个项目之前选过，答案不会变更优\n",
    "            else:\n",
    "                if not h:continue    #如果选上的目前都只有一个,则新添不会使答案更优\n",
    "                else:\n",
    "                    r = heappop(h)\n",
    "                    t = -(nowtype)**2 + (nowtype+1)**2 - r + cost\n",
    "                    res += t\n",
    "                    nowtype += 1\n",
    "                    cnt[types] += 1\n",
    "                    ans = max(ans,res)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x : -x[0])\n",
    "        n = len(items)\n",
    "        ans = 0\n",
    "        cnt = collections.Counter()\n",
    "        sz = 0 \n",
    "        for i in range(k):\n",
    "            ans += items[i][0]\n",
    "            cnt[items[i][1]] += 1\n",
    "            if cnt[items[i][1]] == 1:\n",
    "                sz += 1\n",
    "        ans += sz * sz \n",
    "        Q = []\n",
    "        for p, c in items[:k]:\n",
    "            if cnt[c] > 1:\n",
    "                heapq.heappush(Q, (p, c))\n",
    "\n",
    "        res = ans \n",
    "        for np, nc in items[k:]:\n",
    "            if nc in cnt:\n",
    "                continue \n",
    "            if not Q:\n",
    "                break \n",
    "            p, c = heapq.heappop(Q) \n",
    "            while Q and cnt[c] <= 1:\n",
    "                p, c = heapq.heappop(Q) \n",
    "\n",
    "            if cnt[c] > 1:\n",
    "                ans += np + 1 + 2 * sz - p \n",
    "                sz += 1\n",
    "                cnt[nc] += 1\n",
    "                cnt[c] -= 1\n",
    "                res = max(ans, res)\n",
    "        return res \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "        n=len(items)\n",
    "        items.sort(key=lambda x:-x[0])\n",
    "        cnt=Counter()\n",
    "        tot=0\n",
    "        vis=set()\n",
    "        arr=[]\n",
    "\n",
    "        for i in range(k):\n",
    "            p,c=items[i]\n",
    "            tot+=p\n",
    "            if cnt[c]:\n",
    "                arr.append((p,c))\n",
    "            cnt[c]+=1\n",
    "\n",
    "        res=tot+len(cnt)**2\n",
    "        for i in range(k,n):\n",
    "            # 将重复的存起来，用新来的替换\n",
    "            p,c=items[i]\n",
    "            if cnt[c]==0 and arr:\n",
    "                p1,c1=arr.pop()\n",
    "                tot-=p1\n",
    "                cnt[c1]-=1\n",
    "                if cnt[c1]==0:del cnt[c1]\n",
    "                tot+=p\n",
    "                cnt[c]+=1\n",
    "            res=max(res,tot+len(cnt)**2)\n",
    "\n",
    "        return res\n",
    "            \n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        items.sort(key=lambda x : -x[0])\n",
    "        tot = 0\n",
    "        cnt = Counter()\n",
    "        j = -1\n",
    "        for i in range(k):\n",
    "            tot += items[i][0]\n",
    "            cnt[items[i][1]] += 1\n",
    "            if cnt[items[i][1]] > 1:\n",
    "                j = i\n",
    "        if len(cnt) == k:\n",
    "            return tot + len(cnt) ** 2 \n",
    "        ans = tot + len(cnt) ** 2  \n",
    "        for i in range(k, n):\n",
    "            if items[i][1] in cnt:\n",
    "                continue \n",
    "            tot += items[i][0]\n",
    "            tot -= items[j][0]\n",
    "            cnt[items[j][1]] -= 1\n",
    "            cnt[items[i][1]] = 1\n",
    "            ans = max(ans, tot + len(cnt) ** 2) \n",
    "            j-=1\n",
    "            while j >= 0 and cnt[items[j][1]] == 1:\n",
    "                j -= 1 \n",
    "            if j < 0:\n",
    "                break \n",
    "        return ans \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x:-x[0])\n",
    "        heap = []\n",
    "        ret = 0\n",
    "        dic = {}\n",
    "        for pro,cat in items[:k]:\n",
    "            ret += pro\n",
    "            if cat not in dic:\n",
    "                dic[cat] = pro\n",
    "            else:\n",
    "                heap.append(pro)\n",
    "        heapify(heap)\n",
    "        ans = ret + len(dic)**2\n",
    "        for pro,cat in items[k:]:\n",
    "            if cat in dic:\n",
    "                continue\n",
    "            if not heap:\n",
    "                continue\n",
    "            # if pro - heap[0] + (cur_cat+1)**2 - cur_cat**2 > 0:\n",
    "            dic[cat] = pro\n",
    "            ret += pro - heap[0]\n",
    "            heappop(heap)\n",
    "            ans = max(ans,ret + len(dic.keys())**2)\n",
    "        return ans\n",
    "                \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key = lambda x:x[0], reverse = True)\n",
    "        count_dict = dict()\n",
    "        total_p = 0\n",
    "        for i in range(k):\n",
    "            p, c = items[i]\n",
    "            total_p += p \n",
    "            count_dict[c] = count_dict.get(c, 0) + 1\n",
    "        ret = total_p + (len(count_dict)) ** 2\n",
    "        n = min(len(set(item[1] for item in items)), k)\n",
    "        i = k-1\n",
    "        j = k\n",
    "        while len(count_dict) < n:\n",
    "            while i >= 0:\n",
    "                if count_dict[items[i][1]] > 1:\n",
    "                    total_p -= items[i][0]\n",
    "                    count_dict[items[i][1]] -= 1\n",
    "                    i -= 1\n",
    "                    break\n",
    "                i -= 1\n",
    "            else:\n",
    "                break\n",
    "            while j < len(items):\n",
    "                if items[j][1] not in count_dict:\n",
    "                    total_p += items[j][0]\n",
    "                    count_dict[items[j][1]] = 1\n",
    "                    j += 1\n",
    "                    break\n",
    "                j += 1\n",
    "            else:\n",
    "                break\n",
    "            cur = total_p + (len(count_dict)) ** 2\n",
    "            if cur > ret:\n",
    "                ret = cur\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x:-x[0])\n",
    "        S=0\n",
    "        dic=dict()\n",
    "        for i in range(k):\n",
    "            dic[items[i][1]]=dic.get(items[i][1],0)+1\n",
    "            S+=items[i][0]\n",
    "        S+=len(dic)**2\n",
    "        L=items[:k]\n",
    "        ans=[S]\n",
    "        for i in range(k,len(items)):\n",
    "            if dic.get(items[i][1],0)==0:\n",
    "                while len(L)>0 and dic[L[-1][1]]==1:\n",
    "                    L.pop()\n",
    "                if len(L)==0:\n",
    "                    break\n",
    "                else:\n",
    "                    dic[L[-1][1]]-=1\n",
    "                    dic[items[i][1]]=1\n",
    "                    ans.append(ans[-1]-L[-1][0]+2*len(dic)-1+items[i][0])\n",
    "                    L.pop()\n",
    "        return max( ans )\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        max_profits = sorted(items, key=lambda x: -x[0])\n",
    "        v_sum = 0\n",
    "        c_cnter = collections.Counter()\n",
    "        for p, c in max_profits[:k]:\n",
    "            v_sum += p\n",
    "            c_cnter[c] += 1\n",
    "        c_num = len(c_cnter)\n",
    "        ans = v_sum + c_num ** 2\n",
    "\n",
    "        c_max = collections.defaultdict(int)\n",
    "        for p, c in max_profits[k:]:\n",
    "            if c in c_cnter:\n",
    "                continue\n",
    "            c_max[c] = max(c_max[c], p)\n",
    "        choices = sorted(c_max.values(), reverse=True)\n",
    "\n",
    "        m = len(choices)\n",
    "        i, j = k - 1, 0\n",
    "        while i >= 0 and j < m:\n",
    "            p, c = max_profits[i]\n",
    "            if c_cnter[c] > 1:\n",
    "                c_cnter[c] -= 1\n",
    "                v_sum += choices[j] - p\n",
    "                c_num += 1\n",
    "                ans = max(ans, v_sum + c_num ** 2)\n",
    "                j += 1\n",
    "            i -= 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        items.sort(key = lambda x: -x[0])\n",
    "        d = defaultdict(int) # 不同种类的个数\n",
    "        cnt = 0\n",
    "        cur = 0 \n",
    "        from sortedcontainers import SortedList\n",
    "        repeats = SortedList()\n",
    "        for i in range(k):\n",
    "            cur += items[i][0]\n",
    "            if d[items[i][1]] == 0:\n",
    "                cnt += 1\n",
    "            else:\n",
    "                repeats.add(items[i])\n",
    "            d[items[i][1]] += 1\n",
    "        cur += cnt * cnt\n",
    "        ans = cur\n",
    "        if len(repeats) == 0:\n",
    "            return ans\n",
    "        for i in range(k, n):\n",
    "            if d[items[i][1]] == 0: # 增加新的类型\n",
    "                if len(repeats) == 0:\n",
    "                    return ans\n",
    "                add = (cnt + 1) * (cnt + 1) - cnt * cnt\n",
    "                minItem = repeats[0]\n",
    "                delete = minItem[0] - items[i][0]\n",
    "                # 不一定要求add > delete\n",
    "                # 有频数大于1的就可以换\n",
    "                # 然后换完取最大就行\n",
    "                cur += add - delete\n",
    "                ans = max(ans, cur)\n",
    "                #repeats.add(items[i])\n",
    "                cnt += 1\n",
    "                d[items[i][1]] = 1\n",
    "                d[minItem[1]] -= 1\n",
    "                repeats.discard(minItem)\n",
    "        return ans\n",
    "                    \n",
    "                \n",
    "            \n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# class Solution:\n",
    "#     def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        a = [[] for _ in range(n)]\n",
    "        u = [-1] * n\n",
    "        v = []\n",
    "\n",
    "        for c in items:\n",
    "            a[c[1] - 1].append(c[0])\n",
    "\n",
    "        for i in range(n):\n",
    "            if a[i]:\n",
    "                a[i].sort()\n",
    "                u[i] = a[i].pop()\n",
    "                for c in a[i]:\n",
    "                    v.append(c)\n",
    "\n",
    "        u.sort(reverse=True)\n",
    "        v.sort(reverse=True)\n",
    "        ans = 0\n",
    "        s, t = 0, 0\n",
    "\n",
    "        for i in range(k):\n",
    "            if i < len(v):\n",
    "                s += v[i]\n",
    "\n",
    "        for i in range(k):\n",
    "            if u[i] != -1:\n",
    "                if k - i - 1 < len(v):\n",
    "                    s -= v[k - i - 1]\n",
    "                t += u[i]\n",
    "                ans = max(ans, s + t + (i + 1) * (i + 1))\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "        itemk = []\n",
    "        items = sorted(items,key = lambda x:-x[0])\n",
    "        res0 = 0\n",
    "        itemscounter = Counter([i[1] for i in items[0:k]])\n",
    "        knum = len(itemscounter)\n",
    "        res = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        res0 = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        for i in range(k):\n",
    "            if(itemscounter[items[i][1]]>1):\n",
    "                heapq.heappush(itemk,items[i])\n",
    "        for i in range(k,len(items)):\n",
    "            if(items[i][1] in itemscounter):\n",
    "                continue\n",
    "            if(itemk == []):\n",
    "                break\n",
    "            flag = 0\n",
    "            while(itemk!=[]):\n",
    "                t = heapq.heappop(itemk)\n",
    "                if(itemscounter[t[1]]>1):\n",
    "                    flag = 1\n",
    "                    break\n",
    "            if(flag==1):\n",
    "                itemscounter[t[1]] -= 1\n",
    "                res = res+items[i][0]-t[0]\n",
    "                knum += 1\n",
    "                itemscounter[items[i][1]] = 1\n",
    "                res += 2*knum-1\n",
    "            if(res>res0):\n",
    "                res0 = res\n",
    "        if(res>res0):\n",
    "                res0 = res\n",
    "        \n",
    "        return res0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "        itemk = []\n",
    "        items = sorted(items,key = lambda x:-x[0])\n",
    "        res0 = 0\n",
    "        itemscounter = Counter([i[1] for i in items[0:k]])\n",
    "        knum = len(itemscounter)\n",
    "        res = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        res0 = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        for i in range(k):\n",
    "            if(itemscounter[items[i][1]]>1):\n",
    "                heapq.heappush(itemk,items[i])\n",
    "\n",
    "        for i in range(k,len(items)):\n",
    "            #print(itemk,items[i],knum)\n",
    "            if(items[i][1] in itemscounter):\n",
    "                continue\n",
    "            if(itemk == []):\n",
    "                break\n",
    "            flag = 0\n",
    "            while(itemk!=[]):\n",
    "                t = heapq.heappop(itemk)\n",
    "                if(itemscounter[t[1]]>1):\n",
    "                    flag = 1\n",
    "                    break\n",
    "            if(flag==1):\n",
    "                itemscounter[t[1]] -= 1\n",
    "                res = res+items[i][0]-t[0]\n",
    "                knum += 1\n",
    "                itemscounter[items[i][1]] = 1\n",
    "                res += 2*knum-1\n",
    "            if(res>res0):\n",
    "                res0 = res\n",
    "        if(res>res0):\n",
    "                res0 = res\n",
    "        return res0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "        itemk = []\n",
    "        items = sorted(items,key = lambda x:-x[0])\n",
    "        res0 = 0\n",
    "        itemscounter = Counter([i[1] for i in items[0:k]])\n",
    "        knum = len(itemscounter)\n",
    "        res = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        res0 = sum([i[0] for i in items[0:k]]) + knum**2\n",
    "        for i in range(k):\n",
    "            if(itemscounter[items[i][1]]>1):\n",
    "                heapq.heappush(itemk,items[i])\n",
    "\n",
    "        for i in range(k,len(items)):\n",
    "            #print(itemk,items[i],knum)\n",
    "            if(items[i][1] in itemscounter):\n",
    "                continue\n",
    "            if(itemk == []):\n",
    "                break\n",
    "            flag = 0\n",
    "            while(itemk!=[]):\n",
    "                t = heapq.heappop(itemk)\n",
    "                if(itemscounter[t[1]]>1):\n",
    "                    flag = 1\n",
    "                    break\n",
    "            if(flag==1):\n",
    "                itemscounter[t[1]] -= 1\n",
    "                res = res+items[i][0]-t[0]\n",
    "                knum += 1\n",
    "                itemscounter[items[i][1]] = 1\n",
    "                res += 2*knum-1\n",
    "            if(res>res0):\n",
    "                res0 = res\n",
    "        if(res>res0):\n",
    "                res0 = res\n",
    "        return res0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        if len(items) < k:\n",
    "            return 0\n",
    "        items.sort(key = lambda x: x[0], reverse=True)\n",
    "        catFreq = {}\n",
    "        for i in range(k):\n",
    "            if items[i][1] in catFreq:\n",
    "                catFreq[items[i][1]] += 1\n",
    "            else:\n",
    "                catFreq[items[i][1]] = 1\n",
    "        selectedElements = [i for i in range(0, k)]\n",
    "        ans = len(catFreq) ** 2\n",
    "        for i in range(k):\n",
    "            ans += items[i][0]\n",
    "        tmp = ans\n",
    "        i = k\n",
    "        while i < len(items):\n",
    "            if items[i][1] in catFreq:\n",
    "                i += 1\n",
    "            else:\n",
    "                tailElements = []\n",
    "                while selectedElements and catFreq[items[selectedElements[-1]][1]] == 1:\n",
    "                    tailElements.append(selectedElements.pop())\n",
    "                if not selectedElements:\n",
    "                    break\n",
    "                catFreq[items[selectedElements[-1]][1]] -= 1\n",
    "                tmp = tmp - items[selectedElements[-1]][0] + items[i][0]\n",
    "                tmp = tmp - len(catFreq) ** 2 + (len(catFreq) + 1) ** 2\n",
    "                catFreq[items[i][1]] = 1\n",
    "                ans = max(ans, tmp)\n",
    "                tailElements.reverse()\n",
    "                selectedElements.pop()\n",
    "                selectedElements.extend(tailElements)\n",
    "                selectedElements.append(i)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, g: List[List[int]], k: int) -> int:\n",
    "        c = 0\n",
    "        s = 0\n",
    "        h = 0\n",
    "        g.sort(reverse=True)\n",
    "        a = [0 for _ in g]\n",
    "        b = []\n",
    "        for ii in range(len(g)):\n",
    "            i, j = g[ii]\n",
    "            if a[j - 1] == 0:\n",
    "                c += 1\n",
    "                s += i\n",
    "                h += 1\n",
    "                b.append([i, j, ii])\n",
    "                g[ii][1] = -1\n",
    "                a[j - 1] = 1\n",
    "                if h == k:\n",
    "                    break\n",
    "        if h < k:\n",
    "            for ii in range(len(g)):\n",
    "                i, j = g[ii]\n",
    "                if j == -1:\n",
    "                    continue\n",
    "                if h < k:\n",
    "                    s += i\n",
    "                    h += 1\n",
    "                    g[ii][1] = -1\n",
    "                if h == k:\n",
    "                    break\n",
    "        b.sort()\n",
    "        ans = s + c * c\n",
    "        t = 0\n",
    "        for ii in range(len(g)):\n",
    "            i, j = g[ii]\n",
    "            if j == -1:\n",
    "                continue\n",
    "            if s + i - b[t][0] + (c - 1) * (c - 1) > ans:\n",
    "                ans = s + i - b[t][0] + (c - 1) * (c - 1)\n",
    "                s += i - b[t][0]\n",
    "                t += 1\n",
    "                c -= 1\n",
    "            else:\n",
    "                break\n",
    "        return ans\n",
    "            \n",
    "                \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        ## 刚开始按照利润从大到小排序，选择最大的k个利润\n",
    "        n = len(items)\n",
    "        items.sort(key=lambda x:-x[0])\n",
    "        cnt = Counter([c for _,c in items[:k]])\n",
    "        total_c = len(cnt)\n",
    "        res = s = sum(p for p,_ in items[:k]) + total_c**2 # 初始的最大值\n",
    "        cur = list(range(k)) #选中的k个索引\n",
    "        for j in range(k,n):\n",
    "            ## 判断一下第j个数字能不能加进来\n",
    "            p,c = items[j]\n",
    "            if cnt[c]>0:continue # 出现过，加进来没用\n",
    "            ## 没出现过，分两种。\n",
    "            \n",
    "            # 找到需要更换的数字\n",
    "            for i in range(k-1,-1,-1):\n",
    "                candidate = cur[i]\n",
    "                p1,c1 = items[candidate]\n",
    "                if cnt[c1]==1:continue #一种是要踢到的数字只出现一次，不能换\n",
    "                ## 加进来吧\n",
    "                cur.pop(i)\n",
    "                cur.append(j)\n",
    "                cnt[c1]-=1\n",
    "                cnt[c]+=1\n",
    "                s = s-p1-total_c**2 + (total_c+1)**2 + p\n",
    "                total_c += 1\n",
    "                res = max(res,s)\n",
    "                break\n",
    "        return res\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        li = len(items)\n",
    "        subitem = []\n",
    "        table={}\n",
    "        for i in range(li):\n",
    "            prof = items[i][0]\n",
    "            cate = items[i][1]\n",
    "            if table.get(cate,-1)==-1:\n",
    "                table[cate] = prof\n",
    "            else:\n",
    "                if table[cate]>=prof:\n",
    "                    subitem.append(prof)\n",
    "                else:\n",
    "                    subitem.append(table[cate])\n",
    "                    table[cate] = prof\n",
    "        subitem.sort(reverse=-1)\n",
    "        lsub = len(subitem)\n",
    "        ss_sub = subitem.copy()\n",
    "        for i in range(1,lsub):\n",
    "            ss_sub[i]+=ss_sub[i-1]\n",
    "        ss_sub=[0]+ss_sub\n",
    "        topitem = []\n",
    "        for key in table:\n",
    "            topitem.append(table[key])\n",
    "        topitem.sort(reverse=-1)\n",
    "        ltop = len(topitem)\n",
    "        ss_top = topitem.copy()\n",
    "        for i in range(1,ltop):\n",
    "            ss_top[i]+=ss_top[i-1]\n",
    "        ss_top= [0]+ss_top\n",
    "        ans = 0\n",
    "        for i in range(1,k+1):\n",
    "            if i<=ltop and k-i<=lsub:\n",
    "\n",
    "                ans = max(ans,ss_top[i]+ss_sub[k-i]+i**2)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(reverse = True)\n",
    "        d = set()\n",
    "        cates = set()\n",
    "        ans = 0\n",
    "        firsts = []\n",
    "        for i,[profit,category] in enumerate(items):\n",
    "            if i < k:\n",
    "                ans += profit\n",
    "                cates.add(category)\n",
    "            if category not in d:\n",
    "                firsts.append(i)\n",
    "            d.add(category)\n",
    "        num = len(cates)\n",
    "        ans += num * num\n",
    "        numTotal = min(len(firsts),k)\n",
    "        if num >= numTotal:\n",
    "            return ans\n",
    "        firsts.sort(key = lambda x:items[x][0],reverse = True)\n",
    "        tmp = 0\n",
    "        for x in firsts[:numTotal]:\n",
    "            tmp += items[x][0]\n",
    "        n = numTotal\n",
    "        m = 0\n",
    "        while n < k:\n",
    "            if m not in firsts:\n",
    "                tmp += items[m][0]\n",
    "                n += 1\n",
    "            m += 1\n",
    "        ans = max(ans,tmp + numTotal * numTotal)        \n",
    "        for i in range(numTotal - 1,num, -1):\n",
    "            tmp -= items[firsts[i]][0]\n",
    "            n -= 1\n",
    "            while n < k:\n",
    "                if m not in firsts:\n",
    "                    tmp += items[m][0]\n",
    "                    n += 1\n",
    "                m += 1\n",
    "            ans = max(ans,tmp + i * i)\n",
    "        return ans\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(reverse = True)\n",
    "        cates = set()\n",
    "        first = set()\n",
    "        tmp = 0\n",
    "        for i,[profit,category] in enumerate(items):\n",
    "            if i < k:\n",
    "                if category not in cates:\n",
    "                    first.add(i)\n",
    "                cates.add(category)\n",
    "                tmp += profit\n",
    "                if i == k - 1:\n",
    "                    j = len(cates)\n",
    "                    tmp += j * j\n",
    "                    if j == k:\n",
    "                        return tmp\n",
    "                    ans = tmp\n",
    "                    n = i\n",
    "            else:\n",
    "                if category in cates:\n",
    "                    continue\n",
    "                else:\n",
    "                    cates.add(category)\n",
    "                    while n in first:\n",
    "                        n -= 1\n",
    "                    if n < 0:\n",
    "                        return ans\n",
    "                    else:\n",
    "                        j += 1\n",
    "                        tmp += j * j + profit - (j - 1) * (j - 1) - items[n][0]\n",
    "                        n -= 1\n",
    "                        ans = max(ans,tmp)\n",
    "                        if n < 0:\n",
    "                            return ans\n",
    "        return ans\n",
    "        \n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(reverse = True)\n",
    "        print(items)\n",
    "        cates = set()\n",
    "        first = set()\n",
    "        tmp = 0\n",
    "        for i,[profit,category] in enumerate(items):\n",
    "            if i < k:\n",
    "                if category not in cates:\n",
    "                    first.add(i)\n",
    "                cates.add(category)\n",
    "                tmp += profit\n",
    "                if i == k - 1:\n",
    "                    j = len(cates)\n",
    "                    tmp += j * j\n",
    "                    if j == k:\n",
    "                        return tmp\n",
    "                    ans = tmp\n",
    "                    n = i\n",
    "                    print(tmp,j)\n",
    "            else:\n",
    "                if category in cates:\n",
    "                    continue\n",
    "                else:\n",
    "                    cates.add(category)\n",
    "                    while n in first:\n",
    "                        n -= 1\n",
    "                    if n < 0:\n",
    "                        return ans\n",
    "                    else:\n",
    "                        j += 1\n",
    "                        tmp += j * j + profit - (j - 1) * (j - 1) - items[n][0]\n",
    "                        n -= 1\n",
    "                        ans = max(ans,tmp)\n",
    "                        if n < 0:\n",
    "                            return ans\n",
    "        return ans\n",
    "        \n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x: -x[0])\n",
    "        cnt = Counter()\n",
    "        cur = 0\n",
    "        st = []\n",
    "        for i, (x, y) in enumerate(items[:k]):\n",
    "            cur += x\n",
    "            cnt[y] += 1\n",
    "            st.append((x, y))\n",
    "\n",
    "        tot = len(cnt)\n",
    "        cur += tot ** 2\n",
    "        res = cur\n",
    "        for x, y in items[k:]:\n",
    "            if y not in cnt or cnt[y] == 0:\n",
    "\t\t\t\t\t\t\t\t# 懒删除\n",
    "                while st and cnt[st[-1][1]] == 1:\n",
    "                    nx, ny = st.pop()\n",
    "                # 反悔操作（保证种类越多越好）\n",
    "                if st:\n",
    "                    nx, ny = st.pop()\n",
    "                    cur += x - nx - (tot ** 2) + ((tot + 1) ** 2)\n",
    "                    tot += 1\n",
    "                    cnt[ny] -= 1\n",
    "                    cnt[y] += 1\n",
    "                    res = max(res, cur)\n",
    "                else:\n",
    "                    break\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items = sorted(items, key=lambda x:-x[0])\n",
    "        ans = 0\n",
    "        selects = {}\n",
    "        category = {}\n",
    "        elements = []\n",
    "        \n",
    "        for i in range(len(items)):\n",
    "            item = items[i]\n",
    "            selects[i] = 1\n",
    "            if item[1] in category:\n",
    "                category[item[1]] += 1\n",
    "            else:\n",
    "                category[item[1]] = 1\n",
    "            ans += item[0]\n",
    "            if category[item[1]] >= 2:\n",
    "                elements.append(item)\n",
    "            k -= 1\n",
    "            if k == 0:\n",
    "                break\n",
    "        \n",
    "        #print(ans, category)\n",
    "        ans += len(category) * len(category)\n",
    "        \n",
    "        elements = sorted(elements, key=lambda x:-x[0])\n",
    "        \n",
    "        if len(elements) == 0:\n",
    "            return ans\n",
    "        \n",
    "        total = ans\n",
    "                \n",
    "        for i in range(len(items)):\n",
    "            if len(elements) == 0:\n",
    "                return ans\n",
    "            if i in selects:\n",
    "                continue\n",
    "            if items[i][1] in category:\n",
    "                continue\n",
    "            else:\n",
    "                item = elements[-1]\n",
    "                if item[1] in category and category[item[1]] == 1:\n",
    "                    continue\n",
    "                \n",
    "                category[item[1]] -= 1\n",
    "                item = elements.pop()\n",
    "                \n",
    "                \n",
    "                if items[i][1] in category:\n",
    "                    category[items[i][1]] += 1\n",
    "                else:\n",
    "                    category[items[i][1]] = 1\n",
    "                \n",
    "                # if category[items[i][1]] > 1:\n",
    "                #     elements.append(items[i])\n",
    "                # print(ans, item[0], len(category), items[i][0])\n",
    "                total = total - item[0] + 2 * len(category) + items[i][0] - 1\n",
    "                ans = max(ans, total)\n",
    "                \n",
    "        \n",
    "        return ans\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        dc = {}\n",
    "        it = sorted(items,reverse=True)\n",
    "        n = len(it)\n",
    "        sm, m = 0, 0\n",
    "        lis = []\n",
    "        for i,j in it[:k]:\n",
    "            if j in dc:\n",
    "                lis.append(i)\n",
    "            else:\n",
    "                dc[j] = []\n",
    "            sm += i\n",
    "        lis.sort(reverse=True)\n",
    "        m = len(dc)\n",
    "        mx = sm + m * m\n",
    "        if not lis: return mx\n",
    "        # print(sm, m, lis)\n",
    "        for i,j in it[k:]:\n",
    "            if j not in dc:\n",
    "                dc[j] = []\n",
    "                sm += i - lis.pop()\n",
    "                m += 1\n",
    "                mx = max(mx, sm + m * m)\n",
    "                if not lis: break\n",
    "        return mx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        items.sort(reverse=True)\n",
    "        base = 0\n",
    "        s = defaultdict(int)\n",
    "        heap = []\n",
    "        heapify(heap)\n",
    "        for i in range(k):\n",
    "            pro, cat = items[i]\n",
    "            s[cat] += 1\n",
    "            base += pro\n",
    "            heappush(heap, [pro, cat])\n",
    "            \n",
    "        \n",
    "        \n",
    "        ans = base + len(s) ** 2\n",
    "    \n",
    "        \n",
    "        for i in range(k, n):\n",
    "            pro, cat = items[i]\n",
    "            if s[cat] == 0:\n",
    "                success = 0\n",
    "                pp = -1\n",
    "                while len(heap):\n",
    "                    p, c = heappop(heap)\n",
    "                    if s[c] != 1:\n",
    "                        s[c] -= 1\n",
    "                        success = 1\n",
    "                        pp = p\n",
    "                        break\n",
    "                \n",
    "                if success:\n",
    "                    s[cat] = 1\n",
    "                    base -= pp\n",
    "                    base += pro\n",
    "                    ans = max(ans, base + len(s) ** 2)\n",
    "                \n",
    "        return ans\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        mmap = {}\n",
    "        N = len(items)\n",
    "\n",
    "        for p, c in items:\n",
    "            if c not in mmap:\n",
    "                mmap[c] = [p]\n",
    "            else:\n",
    "                bisect.insort(mmap[c], p)\n",
    "\n",
    "        n = len(mmap)\n",
    "        items.sort()\n",
    "\n",
    "        one_q = []\n",
    "        more_q = []\n",
    "        while N > k:\n",
    "            index = None\n",
    "            minv = None\n",
    "            cgry = None\n",
    "            for i in range(len(items)):\n",
    "                v, c = items[i]\n",
    "                if index is None:\n",
    "                    if len(mmap[c]) > 1 or k < n:\n",
    "                        mmap[c].pop(0)\n",
    "                        if len(mmap[c]) == 0:\n",
    "                            del mmap[c]\n",
    "                            n -= 1\n",
    "                        items.pop(i)\n",
    "                        break\n",
    "                    else:\n",
    "                        index = i\n",
    "                        minv = v\n",
    "                        cgry = c\n",
    "                        continue\n",
    "                else:\n",
    "                    if len(mmap[c]) == 1:\n",
    "                        continue\n",
    "                    if n ** 2 - (n - 1) ** 2 > v - minv:\n",
    "                        mmap[c].pop(0)\n",
    "                        items.pop(i)\n",
    "                    else:\n",
    "                        mmap[cgry].pop()\n",
    "                        del mmap[cgry]\n",
    "                        n -= 1\n",
    "                        items.pop(index)\n",
    "                    break\n",
    "            N -= 1\n",
    "\n",
    "        result = 0\n",
    "        for c, pq in mmap.items():\n",
    "            result += sum(pq)\n",
    "        result += n ** 2\n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x: -x[0])\n",
    "        hst = Counter()\n",
    "        total_profit = 0\n",
    "        h = []\n",
    "        for p, c in items[:k]:\n",
    "            total_profit += p\n",
    "            hst[c] += 1\n",
    "            heapq.heappush(h, [p, c])\n",
    "        distinct_category = len(hst.keys())\n",
    "        rst = total_profit + distinct_category ** 2\n",
    "        for p, c in items[k:]:\n",
    "            if distinct_category < k and c not in hst:\n",
    "                while True:\n",
    "                    x, y = heapq.heappop(h)\n",
    "                    if hst[y] == 1:\n",
    "                        continue\n",
    "                    else:\n",
    "                        hst[y] -= 1\n",
    "                        break\n",
    "                total_profit += p - x\n",
    "                distinct_category += 1\n",
    "                hst[c] = 1\n",
    "                rst = max(rst, total_profit + distinct_category ** 2)\n",
    "        return rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        from collections import defaultdict\n",
    "        from sortedcontainers import SortedList\n",
    "\n",
    "        items.sort(reverse=True)\n",
    "\n",
    "        cate = Counter()\n",
    "        pick = SortedList()\n",
    "        count = 0\n",
    "        for p, c in items[:k]:\n",
    "            cate[c] += 1\n",
    "            pick.add([p, c])\n",
    "            count += p\n",
    "        ans = count + len(cate) ** 2\n",
    "\n",
    "        for p, c in items[k:]:\n",
    "            if c in cate:\n",
    "                continue\n",
    "            for pp, cc in pick:\n",
    "                if cate[cc] > 1:\n",
    "                    pick.remove([pp, cc])\n",
    "                    pick.add([p, c])\n",
    "                    count -= pp\n",
    "                    count += p\n",
    "                    cate[cc] -= 1\n",
    "                    cate[c] += 1\n",
    "                    break\n",
    "            ans = max(ans, count + len(cate) ** 2)\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        A = sorted(items,reverse=True)\n",
    "        counter = Counter()\n",
    "        que = []\n",
    "        sm_p = 0\n",
    "        for p,c in A[:k]:\n",
    "            que.append([p,c])\n",
    "            counter[c]+=1\n",
    "            sm_p += p\n",
    "        res = sm_p + len(counter)**2\n",
    "        for p,c in A[k:]:\n",
    "            if c not in counter:\n",
    "                while que:\n",
    "                    p2,c2 = que.pop()\n",
    "                    if counter[c2]>=2:\n",
    "                        counter[c2]-=1\n",
    "                        sm_p -= p2\n",
    "                        break\n",
    "                if not que:\n",
    "                    break\n",
    "                sm_p += p\n",
    "                counter[c]+=1\n",
    "                res = max(res, sm_p+len(counter)**2)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        \n",
    "        \n",
    "        \n",
    "        # my solution ... \n",
    "        #  time: O()\n",
    "        # space: O()\n",
    "        \n",
    "        items.sort(key = lambda x: -x[0])\n",
    "        psum = 0\n",
    "        c2p = {}\n",
    "        res, pq = 0, []\n",
    "        for index, (profit, category) in enumerate(items):\n",
    "            if index < k:\n",
    "                psum += profit\n",
    "                if category not in c2p:\n",
    "                    c2p[category] = []\n",
    "                c2p[category].append(profit)\n",
    "                if index == k-1:\n",
    "                    res = psum + len(c2p) ** 2\n",
    "                    for c in c2p:\n",
    "                        if len(c2p[c]) > 1:\n",
    "                            for p in c2p[c]:\n",
    "                                pq.append((p,c))\n",
    "                    heapq.heapify(pq)\n",
    "            elif category not in c2p:\n",
    "                while pq and len(c2p[pq[0][1]]) == 1:\n",
    "                    heapq.heappop(pq)\n",
    "                if not pq:\n",
    "                    break\n",
    "                p, c = heapq.heappop(pq)\n",
    "                psum += profit - p\n",
    "                c2p[c].pop()\n",
    "                c2p[category] = [profit]\n",
    "                res = max(res, psum + len(c2p) ** 2)\n",
    "        return res\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "\n",
    "        items.sort(reverse=True)\n",
    "        ans = cnt = pre = tp = 0\n",
    "        dct = defaultdict(list)\n",
    "        stack = []\n",
    "        for p, c in items:\n",
    "            if cnt == k:\n",
    "                while stack and len(dct[stack[-1]]) == 1:\n",
    "                    stack.pop()\n",
    "                if not stack:\n",
    "                    break\n",
    "                pre -= dct[stack.pop()].pop()\n",
    "            else:\n",
    "                cnt += 1\n",
    "            pre += p\n",
    "            dct[c].append(p)\n",
    "            if len(dct[c]) == 1:\n",
    "                tp += 1\n",
    "            stack.append(c)\n",
    "            if pre + tp * tp > ans:\n",
    "                ans = pre + tp * tp\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from heapq import heapify, heappop, heappush\r\n",
    "class Solution:\r\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\r\n",
    "        n = len(items)\r\n",
    "        items.sort(reverse=True)\r\n",
    "        # catagory -> profits\r\n",
    "        select: dict[int, list[int]] = dict()\r\n",
    "        # profit, catagory\r\n",
    "        finder: list[tuple[int, int]] = list()\r\n",
    "        heapify(finder)\r\n",
    "        # select the top profitable k items\r\n",
    "        for i in range(k):\r\n",
    "            if items[i][1] not in select: select[items[i][1]] = list()\r\n",
    "            select[items[i][1]].append(items[i][0])\r\n",
    "        # add all profits, who has more than 1 elements in their catagory to heap\r\n",
    "        for cat, pfts in select.items():\r\n",
    "            for pft in pfts[1:]:\r\n",
    "                heappush(finder, (pft, cat))\r\n",
    "        # try rest of profits\r\n",
    "        s = sum(sum(v) for v in select.values()) + len(select) * len(select)\r\n",
    "        res = s\r\n",
    "        for i in range(k, n):\r\n",
    "            # if an item catagory already select, it will not increase elegance, as previous profits are larger\r\n",
    "            if items[i][1] in select: continue\r\n",
    "            # select this item may increase elegance, only if it increase distinct catagory, and the increase covers the decrease\r\n",
    "            # in unselecting an item with larger profit\r\n",
    "            if not finder: break\r\n",
    "            s += (len(select) + 1) * (len(select) + 1) + items[i][0] - len(select) * len(select) - finder[0][0]\r\n",
    "            res = max(res, s)\r\n",
    "            select[heappop(finder)[1]].pop()\r\n",
    "            select[items[i][1]] = [items[i][0]]\r\n",
    "\r\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        n = len(items)\n",
    "        \n",
    "        # 最大堆\n",
    "        left = []\n",
    "        \n",
    "        cnt = Counter()\n",
    "        for num,t in items:\n",
    "            heappush(left,(-num,t))\n",
    "        \n",
    "        # 最小堆\n",
    "        right = []\n",
    "        \n",
    "        i = 0\n",
    "        ss = 0\n",
    "        while i < k:\n",
    "            num,t = heappop(left)\n",
    "            num = -num\n",
    "            cnt[t] += 1\n",
    "            ss += num\n",
    "            i += 1\n",
    "            heappush(right,(num,t))\n",
    "        \n",
    "        \n",
    "        res = ss + len(cnt) * len(cnt)\n",
    "\n",
    "        # 最多k个不同项目\n",
    "        while len(cnt) < k:\n",
    "            # 如果要移除的项目只有一个，肯定不能移除，要保留至少一个\n",
    "            while right and cnt[right[0][1]] == 1:\n",
    "                heappop(right)\n",
    "            \n",
    "            if not right:\n",
    "                break\n",
    "            \n",
    "            # 移除\n",
    "            num,t = heappop(right)\n",
    "            cnt[t] -= 1\n",
    "            # 更新 ss\n",
    "            ss -= num\n",
    "            \n",
    "            if not left:\n",
    "                break\n",
    "            \n",
    "            # 添加新项目\n",
    "            num,t = heappop(left)\n",
    "            num = -num\n",
    "            heappush(right,(num,t))\n",
    "            cnt[t] += 1\n",
    "            # 更新 ss\n",
    "            ss += num\n",
    "\n",
    "            # 更新最大优雅度\n",
    "            res = max(res,ss+len(cnt)*len(cnt))\n",
    "        \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items = sorted(items, key=lambda x:x[0], reverse=True)\n",
    "        print(items)\n",
    "        selected = defaultdict(list)\n",
    "        q = []\n",
    "        cate_nums = 0\n",
    "        total_profit = 0\n",
    "        for i in range(k):\n",
    "            profit, cate = items[i]\n",
    "            selected[cate].append(profit)\n",
    "            total_profit += profit\n",
    "            if len(selected[cate])>1:\n",
    "                heapq.heappush(q, (profit, cate))\n",
    "            else:\n",
    "                cate_nums += 1\n",
    "        ans = total_profit + cate_nums**2\n",
    "        for i in range(k, len(items)):\n",
    "            if len(selected.keys())==k:\n",
    "                break\n",
    "            profit, cate = items[i]\n",
    "            if cate in selected:\n",
    "                continue\n",
    "            else:\n",
    "                profit_drop, cate_drop = q[0]\n",
    "                # if (profit-profit_drop+(cate_nums+1)**2-cate_nums**2)>0:\n",
    "                selected[cate].append(profit)\n",
    "                heapq.heappop(q)\n",
    "                selected[cate_drop].pop()\n",
    "                total_profit = total_profit + profit - profit_drop\n",
    "                cate_nums += 1\n",
    "                ans = max(ans, total_profit+cate_nums**2)\n",
    "        \n",
    "        # print(selected)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# -*- coding: utf-8 -*-\n",
    "from typing import List, Tuple, Optional\n",
    "from collections import deque, Counter\n",
    "from queue import PriorityQueue\n",
    "import math\n",
    "from functools import lru_cache\n",
    "from sortedcontainers import SortedDict, SortedSet\n",
    "import random\n",
    "import copy\n",
    "import sys\n",
    "sys.setrecursionlimit(9999999)\n",
    "\n",
    "MOD = 10**9 + 7\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(reverse = True)\n",
    "\n",
    "        mm = {}\n",
    "        tot = 0\n",
    "        for v, g in items[:k]:\n",
    "            if g not in mm:\n",
    "                mm[g] = []\n",
    "            mm[g].append(v)\n",
    "            tot += v\n",
    "\n",
    "        que2 = PriorityQueue()      # 至少有两个数值的组 (组中最小值，分组号)\n",
    "        cat = len(mm)\n",
    "        tot += cat**2\n",
    "\n",
    "        for g, arr in mm.items():\n",
    "            if len(arr) >= 2:\n",
    "                que2.put((arr[-1], g))\n",
    "\n",
    "        ans = tot\n",
    "        for v, g in items[k:]:\n",
    "            if g in mm:\n",
    "                continue\n",
    "\n",
    "            if que2.qsize() == 0:\n",
    "                continue\n",
    "\n",
    "            mn_val, mn_g = que2.queue[0]\n",
    "            arr = mm[mn_g]\n",
    "\n",
    "            delta = (-mn_val) + v + ((cat+1)**2 - cat**2)\n",
    "            arr.pop(-1)\n",
    "            que2.get()\n",
    "            if len(arr) >= 2:\n",
    "                que2.put( (arr[-1], mn_g) )\n",
    "\n",
    "            mm[g] = [v]\n",
    "            tot += delta\n",
    "            cat += 1\n",
    "\n",
    "            ans = max(ans, tot)\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(reverse = True)\n",
    "        n = len(items)\n",
    "        more = []\n",
    "        had = collections.defaultdict(list)\n",
    "        ans = 0\n",
    "        for p, c in items[:k]:\n",
    "            ans += p\n",
    "            heapq.heappush(had[c], p)\n",
    "        ans += len(had) * len(had)\n",
    "\n",
    "        for c in had:\n",
    "            if len(had[c]) > 1:\n",
    "                heapq.heappush(more, (had[c][0], c))\n",
    "        \n",
    "        sums = ans\n",
    "\n",
    "        for p, c in items[k:]:\n",
    "            if c in had:\n",
    "                continue\n",
    "            if not more:\n",
    "                return ans    \n",
    "            sums -= more[0][0]\n",
    "            sums += p\n",
    "            d = len(had)\n",
    "            sums -= d * d\n",
    "            sums += (d + 1) * (d + 1)\n",
    "            if len(had[more[0][1]]) == 2:\n",
    "                heapq.heappop(more)\n",
    "            else:\n",
    "\n",
    "                _, temp = heapq.heappop(more)\n",
    "                heapq.heappop(had[temp])\n",
    "                heapq.heappush(more, (had[temp][0], temp))    \n",
    "            ans = max(ans, sums)\n",
    "            had[c].append(p)\n",
    "\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        # start with the highest costs\n",
    "        # then try out distinct categories\n",
    "        \n",
    "        items.sort(reverse=True)\n",
    "        \n",
    "        taken = items[:k]\n",
    "        remaining = items[k:]\n",
    "        \n",
    "        g = defaultdict(list)\n",
    "        for a,b in taken:\n",
    "            g[b].append(a)\n",
    "\n",
    "        outgoing = []\n",
    "        for k in g.keys():\n",
    "            g[k].sort()\n",
    "            outgoing.extend(g[k][:-1])\n",
    "        outgoing.sort()\n",
    "        \n",
    "        banned_keys = set(g.keys())\n",
    "        f = defaultdict(int)\n",
    "        for a,b in remaining:\n",
    "            if b not in banned_keys:\n",
    "                f[b] = max(f[b], a)\n",
    "        \n",
    "        incoming = list(f.values())\n",
    "        \n",
    "        # print(outgoing)\n",
    "        # print(incoming)\n",
    "                    \n",
    "        curgroups = len(g)\n",
    "        cursum = sum(x[0] for x in taken)\n",
    "        curval = cursum + curgroups ** 2\n",
    "        maxval = curval\n",
    "        # print(cursum, curgroups, curval)\n",
    "        \n",
    "        for x,y in zip(outgoing, incoming):\n",
    "            cursum -= x\n",
    "            cursum += y\n",
    "            curgroups += 1\n",
    "            curval = cursum + curgroups ** 2\n",
    "            maxval = max(maxval, curval)\n",
    "            # print(curval)\n",
    "            \n",
    "        # print()\n",
    "        return maxval\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:\n",
    "        items.sort(key=lambda x : -x[0])\n",
    "        n = len(items)\n",
    "        ans = 0\n",
    "        cnt = collections.Counter()\n",
    "        sz = 0 \n",
    "        for i in range(k):\n",
    "            ans += items[i][0]\n",
    "            cnt[items[i][1]] += 1\n",
    "            if cnt[items[i][1]] == 1:\n",
    "                sz += 1\n",
    "        ans += sz * sz \n",
    "        Q = []\n",
    "        for p, c in items[:k]:\n",
    "            if cnt[c] > 1:\n",
    "                heapq.heappush(Q, (p, c))\n",
    "        print(ans)\n",
    "        print(Q)\n",
    "        print(cnt)\n",
    "        res = ans \n",
    "        for np, nc in items[k:]:\n",
    "            if nc in cnt:\n",
    "                continue \n",
    "            if not Q:\n",
    "                break \n",
    "            p, c = heapq.heappop(Q) \n",
    "            while Q and cnt[c] <= 1:\n",
    "                p, c = heapq.heappop(Q) \n",
    "\n",
    "            if cnt[c] > 1:\n",
    "                ans += np + 1 + 2 * sz - p \n",
    "                sz += 1\n",
    "                cnt[nc] += 1\n",
    "                cnt[c] -= 1\n",
    "                res = max(ans, res)\n",
    "        return res \n",
    "            \n",
    "        "
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
